#include "crypt.h"

static mp_digit prime_tab[] = {
	0x0002, 0x0003, 0x0005, 0x0007, 0x000B, 0x000D, 0x0011, 0x0013, 
	0x0017, 0x001D, 0x001F, 0x0025, 0x0029, 0x002B, 0x002F, 0x0035, 
	0x003B, 0x003D, 0x0043, 0x0047, 0x0049, 0x004F, 0x0053, 0x0059, 
	0x0061, 0x0065, 0x0067, 0x006B, 0x006D, 0x0071, 0x007F, 0x0083, 
	0x0089, 0x008B, 0x0095, 0x0097, 0x009D, 0x00A3, 0x00A7, 0x00AD, 
	0x00B3, 0x00B5, 0x00BF, 0x00C1, 0x00C5, 0x00C7, 0x00D3, 0x00DF, 
	0x00E3, 0x00E5, 0x00E9, 0x00EF, 0x00F1, 0x00FB, 0x0101, 0x0107, 
	0x010D, 0x010F, 0x0115, 0x0119, 0x011B, 0x0125, 0x0133, 0x0137, 

	0x0139, 0x013D, 0x014B, 0x0151, 0x015B, 0x015D, 0x0161, 0x0167, 
	0x016F, 0x0175, 0x017B, 0x017F, 0x0185, 0x018D, 0x0191, 0x0199, 
	0x01A3, 0x01A5, 0x01AF, 0x01B1, 0x01B7, 0x01BB, 0x01C1, 0x01C9, 
	0x01CD, 0x01CF, 0x01D3, 0x01DF, 0x01E7, 0x01EB, 0x01F3, 0x01F7, 
	0x01FD, 0x0209, 0x020B, 0x021D, 0x0223, 0x022D, 0x0233, 0x0239, 
	0x023B, 0x0241, 0x024B, 0x0251, 0x0257, 0x0259, 0x025F, 0x0265, 
	0x0269, 0x026B, 0x0277, 0x0281, 0x0283, 0x0287, 0x028D, 0x0293, 
	0x0295, 0x02A1, 0x02A5, 0x02AB, 0x02B3, 0x02BD, 0x02C5, 0x02CF, 

	0x02D7, 0x02DD, 0x02E3, 0x02E7, 0x02EF, 0x02F5, 0x02F9, 0x0301, 
	0x0305, 0x0313, 0x031D, 0x0329, 0x032B, 0x0335, 0x0337, 0x033B, 
	0x033D, 0x0347, 0x0355, 0x0359, 0x035B, 0x035F, 0x036D, 0x0371, 
	0x0373, 0x0377, 0x038B, 0x038F, 0x0397, 0x03A1, 0x03A9, 0x03AD, 
	0x03B3, 0x03B9, 0x03C7, 0x03CB, 0x03D1, 0x03D7, 0x03DF, 0x03E5, 
	0x03F1, 0x03F5, 0x03FB, 0x03FD, 0x0407, 0x0409, 0x040F, 0x0419, 
	0x041B, 0x0425, 0x0427, 0x042D, 0x043F, 0x0443, 0x0445, 0x0449, 
	0x044F, 0x0455, 0x045D, 0x0463, 0x0469, 0x047F, 0x0481, 0x048B, 

	0x0493, 0x049D, 0x04A3, 0x04A9, 0x04B1, 0x04BD, 0x04C1, 0x04C7, 
	0x04CD, 0x04CF, 0x04D5, 0x04E1, 0x04EB, 0x04FD, 0x04FF, 0x0503, 
	0x0509, 0x050B, 0x0511, 0x0515, 0x0517, 0x051B, 0x0527, 0x0529, 
	0x052F, 0x0551, 0x0557, 0x055D, 0x0565, 0x0577, 0x0581, 0x058F, 
	0x0593, 0x0595, 0x0599, 0x059F, 0x05A7, 0x05AB, 0x05AD, 0x05B3, 
	0x05BF, 0x05C9, 0x05CB, 0x05CF, 0x05D1, 0x05D5, 0x05DB, 0x05E7, 
	0x05F3, 0x05FB, 0x0607, 0x060D, 0x0611, 0x0617, 0x061F, 0x0623,
	0x062B, 0x062F, 0x063D, 0x0641, 0x0647, 0x0649, 0x064D, 0x0653
};

/* figures out if a number is prime (MR test) */
int is_prime(mp_int *N)
{
    long x, s, j;
    int res;
    mp_int n1, a, y, r;
    mp_digit d;

    /* divisible by any of the 256 first primes? */
    for (x = 0; x < 256; x++) {
        mp_mod_d(N, prime_tab[x], &d);
        if (d == 0) {
            return 0;
        }
    }

    /* init variables */
    mp_init_copy(&n1, N); mp_sub_d(&n1, 1, &n1);    /* n1 = N - 1 */
    mp_init_copy(&r, &n1);                          /* r  = N - 1*/
    mp_init(&a);
    mp_init(&y);

    /* find s such that N = (2^s)r */
    s = 0;
    while (mp_iseven(&r)) {
        ++s;
        mp_div_2(&r, &r);
    }
    if (mp_cmp_d(N, 2) == 0)  { res = 1; goto done; }
    if (mp_cmp_d(&r, 0) == 0) { res = 0; goto done; }

    for (x = 0; x < 16; x++) {
        /* choose a */
        mp_set(&a, prime_tab[x]);

        /* compute y = a^r mod n */
        mp_exptmod(&a, &r, N, &y);

        /* (y != 1) AND (y != N-1) */
        if (mp_cmp_d(&y, 1) && mp_cmp(&y, &n1)) {
            for (j = 1; j <= (s-1); j++) {
                /* y = n-1 break */
                if (mp_cmp(&y, &n1) == 0) { break; }

                /* y = y^2 mod N */
                mp_sqrmod(&y, N, &y);

                /* if y == 1 return false */
                if (mp_cmp_d(&y, 1) == 0) { res = 0; goto done; }
            }
            if (mp_cmp(&y, &n1)) { res = 0; goto done; }
        }
    }
    res = 1;
done:
    mp_clear(&a);
    mp_clear(&y);
    mp_clear(&n1);
    mp_clear(&r);
    return res;
}

int rand_prime(mp_int *N, int len, union prng_state *prng, int wprng)
{
   unsigned char buf[4096];
   int x;

   if (len <= 0 || len > 4095) { crypt_error = "Invalid size for rand_prime()."; return CRYPT_ERROR; }
   if (wprng == -1)            { crypt_error = "Invalid prng given to rand_prime()."; return CRYPT_ERROR; }

   buf[0] = 0;
   do {
      x = prng_descriptor[wprng].read(buf+1, len, prng);
      if (x != len) { crypt_error = "Not enough bytes from PRNG in rand_prime()."; return CRYPT_ERROR; }
      buf[1] |= 0x80;
      buf[1+len] |= 0x01;
      mp_read_raw(N, buf, 1+len);
   } while (!is_prime(N));
   memset(buf, 0, 4096);
   return CRYPT_OK;
}
      
